source("/scratch/jc227089/evo-dispersal/evostoch/evostochFunctions.R")
#source("~/evo-dispersal/evostoch/evostochFunctions.R")
setwd("/scratch/jc227089/evostochData/")
#setwd("~/Dropbox/Papers/Submitted/Evolutionary stochasticity/Data/")

fname<-"evostoch"
flist<-list.files(pattern=fname)
flist<-flist[!grepl(pattern="concat", flist)]
flist<-flist[!grepl(pattern="Allee", flist)]

nreps<-20

concat<-c()
#poplist<-vector(mode="list", length=nreps*length(flist))
for (ii in 1:length(flist)) {
	load(flist[ii])
	print(length(out))
	for (jj in 1:length(out)){
		if (is.null(out[[jj]]$pop)) next
		#poplist[[(ii-1)*nreps+(jj)]]<-out[[jj]]$pop
		pars<-rep(unlist(out[[jj]]$parameters), each=2)
		pars<-matrix(pars, nrow=2)
		pars<-cbind(matrix(rep(c(pop=ii, rep=jj), each=2), nrow=2), pars, out[[jj]]$out)
		concat<-rbind(concat, pars)
	}
}
colnames(concat)<-c("pop", "rep", names(unlist(out[[jj]]$parameters)), 
		colnames(out[[jj]]$out))
colnames(concat)[16]<-'xlim'
concat<-cbind(concat, evol=concat[,"h2H"]>0, abs.xlim=abs(concat[,"xlim"]))

concat2<-concat
load(file="concat_evostoch_varR0.RData")
concat2<-rbind(concat2, concat)
concat2<-cbind(concat2, Nstar=(concat2[,"R0"]-1)/concat2[,"a"])

save(concat2, file="concat_evostoch_alphawalker.RData")
#load(file="concat_evostoch_alphawalker.RData")


c.df<-as.data.frame(concat2)
rm(concat, concat2)

# To compare changes in dist and variance subsetted to match 2D case
c.ss<-subset(c.df, (c.df$R0<=10 & c.df$Nstar==10))
meanDist<-tapply(c.ss$abs.xlim, c.ss$evol, mean)
meanDist[2]/meanDist[1]
varDist<-tapply(c.ss$abs.xlim, c.ss$evol, var)
varDist[2]/varDist[1]


var.summ<-c()
var.summ.raw<-c()
for (jj in 1:nlevels(as.factor(c.df$Nstar))){
	for (ii in 1:nlevels(as.factor(c.df$R0))){
		c.tmp<-subset(c.df, as.numeric(as.factor(c.df$R0))==ii & as.numeric(as.factor(c.df$Nstar))==jj)
		c.no.evol<-subset(c.tmp, c.tmp$evol==0)
		c.evol<-subset(c.tmp, c.tmp$evol==1)
		
		dem.st<-summary(aov(abs.xlim~1, data=c.no.evol))[[1]]$'Sum Sq' #pure demographic stochasticity
		evol.st<-summary(aov(abs.xlim~Error(pop), data=c.evol)) # partitioned into within and between pop (between=founder, within = evol+demographic stoch)
		evol.st.win<-evol.st$'Error: Within'[[1]]$'Sum Sq' #evol+dem st
		evol.st.bw<-evol.st$'Error: pop'[[1]]$'Sum Sq' #initial founder event
	
		#proportion attributable to each effect
		dem.st.prop<-dem.st/(evol.st.win+evol.st.bw)
		found.st.prop<-evol.st.bw/(evol.st.win+evol.st.bw)
		evol.st.prop<-(evol.st.win-dem.st)/(evol.st.win+evol.st.bw)
		
		var.summ.raw<-rbind(var.summ.raw, c(R0=as.numeric(levels(as.factor(c.df$R0)))[ii], Nstar=as.numeric(levels(as.factor(c.df$Nstar)))[jj], round(c(demog=dem.st, founder=evol.st.bw, evoln=evol.st.win), 2)))
		var.summ<-rbind(var.summ, c(R0=as.numeric(levels(as.factor(c.df$R0)))[ii], Nstar=as.numeric(levels(as.factor(c.df$Nstar)))[jj], round(c(demog=dem.st.prop, founder=found.st.prop, evoln=evol.st.prop)*100, 2)))
	}
}

var.summ[,3:5]<-var.summ[,3:5]/100
var.summ.raw[,3:5]<-var.summ.raw[,3:5]/400
mean.dist<-tapply(c.df$abs.xlim, list(c.df$R0, c.df$Nstar, c.df$evol), mean)
sd.dist<-tapply(c.df$abs.xlim, list(c.df$R0, c.df$Nstar, c.df$evol), sd)
var.evol<-tapply(c.df$abs.xlim, c.df$evol, var)
####################################################
### Figures ###

lev.ns<-levels(as.factor(var.summ[,"Nstar"]))
fig.loc<-"~/Dropbox/Papers/Submitted/Evolutionary stochasticity/Figures/"
pdf(paste(fig.loc, "alphawalker_panel.pdf", sep=""), height=12, width=4)
	par(mar=c(5,5,1,1), cex.lab=1.4, mfrow=c(3,1))
	X<-as.numeric(dimnames(mean.dist)[[1]])
	matplot(x=X,
		y=cbind(mean.dist[,,2]+sd.dist[,,2], mean.dist[,,2]-sd.dist[,,2]),
		type="n",
		ylab='Realised spread distance',
		xlab=""
		)
		
	
	matplot(x=X, 
		y=mean.dist[,3,2],
		ylim<-range(mean.dist+sd.dist), 
		type="l", 
		bty="l",
		lwd=3,
		lty=1,
		col=1,
		xlab="",#quote(italic(R[max])), 
		add=TRUE
		)
		
	X2<-c(X, rev(X))
	Y2<-c(mean.dist[,3,2]+sd.dist[,3,2], rev(mean.dist[,3,2]-sd.dist[,3,2]))
	polygon(x=X2, 
		y=Y2,
		col=makeTransparent(1),
		border=NA)
		
	matplot(x=X, 
		y=mean.dist[,3,1],
		#ylim<-range(mean.dist), 
		type="l", 
		bty="l",
		lwd=3,
		lty=1,
		col="grey70",
		add=TRUE)	
	legend('topleft', legend=c('Evolution', 'No evolution'), col=c('black', 'grey70'), lty=1, lwd=3, bty='n' )
	
	X2<-c(X, rev(X))
	Y2<-c(mean.dist[,3,1]+sd.dist[,3,1], rev(mean.dist[,3,1]-sd.dist[,3,1]))
	polygon(x=X2, 
		y=Y2,
		col=makeTransparent("grey70"),
		border=NA)
	
	
lev.ns<-levels(as.factor(var.summ.raw[,"Nstar"]))


	matplot(x=var.summ.raw[,1], y=var.summ.raw[,3:5], type="n", bty="l",
		xlab="", ylab='Total variance in spread distance')
	for (nn in 1:length(lev.ns)){
		ss<-as.factor(var.summ.raw[,"Nstar"])==lev.ns[nn]
		for (ii in 3:5){
			lines(var.summ.raw[ss, "R0"], var.summ.raw[ss, ii], lty=ii-2, lwd=nn)
		}
	}	
	
	legend('topleft', legend=c('Demographic stochasticity', 'Initial founder event', 'Evolutionary stochasticity'), col=1, lty=1:3, lwd=3, bty='n' )
	
lev.ns<-levels(as.factor(var.summ[,"Nstar"]))	
	matplot(x=var.summ[,1], 
		y=var.summ[,3:5], 
		type="n", 
		bty="l",
		xlab=quote(italic(R[max])), 
		ylab='Proportion of variance in spread distance')
	for (nn in 1:length(lev.ns)){
		ss<-as.factor(var.summ[,"Nstar"])==lev.ns[nn]
		for (ii in 3:5){
			lines(var.summ[ss, "R0"], var.summ[ss, ii], lty=ii-2, lwd=nn)
		}
	}	
	
	legend('left', legend=c('Demographic stochasticity', 'Initial founder event', 'Evolutionary stochasticity'), col=1, lty=1:3, lwd=3, bty='n' )



dev.off()